#pragma once
#include <memory>
#include "OsqpSolver/OsqpSolver.h"
#include "discrete_linear_system.h"
#define NULL_MAT (Eigen::MatrixXd(0, 0))
#define SAFE_CHEAK
namespace LinearMpc
{
class Mpc
{

public:
	Mpc(std::shared_ptr<DiscreteLinearSys> dis_sys_ptr,
		unsigned int mpc_windows,
		unsigned int max_xu_constraint,
		unsigned int max_delta_x,
		unsigned int max_delta_u,
		unsigned int max_final_constraint,
		unsigned int max_xu_ae = 0,
		unsigned int max_xu_slack_ineq = 0,
		unsigned int max_x_final_ae = 0,
		unsigned int max_x_final_slack_ineq = 0,
		unsigned int max_delta_u_ae = 0,
		unsigned int max_delta_u_slack_ineq = 0,
		unsigned int max_delta_x_ae = 0,
		unsigned int max_delta_x_slack_ineq = 0);

	int set_xu_affine(const std::vector<double> x_index, const std::vector<double> x_factor,
					  const std::vector<double> u_index, const std::vector<double> u_factor,
					  int constraint_index, double min, double max);
	int set_xu_bound(int constraint_index, double min, double max);
	int set_delta_x(unsigned int constraint_index, unsigned int x_index, double min, double max);
	int set_delta_x_bound(unsigned int constraint_index, unsigned int x_index, double min, double max);
	int set_delta_u(unsigned int constraint_index, unsigned int u_index, double min, double max);
	int set_delta_u_bound(unsigned int constraint_index, unsigned int u_index, double min, double max);

	int set_final_constraint(int constraint_index, const std::vector<double> x_index, const std::vector<double> x_factor,
							 double min, double max);
	int set_final_bound(int constraint_index, double min, double max);
	int set_xu_ae(const std::vector<double> x_index, const std::vector<double> x_factor,
				  const std::vector<double> u_index, const std::vector<double> u_factor,
				  int constraint_index, double offset = 0.0);
	bool set_xu_ae_offset(int constraint_index, double offset);
	int set_xu_slack_ineq(const std::vector<double> x_index, const std::vector<double> x_factor,
						  const std::vector<double> u_index, const std::vector<double> u_factor,
						  int constraint_index, double min, double max);
	int set_xu_slack_bound(int constraint_index, double min, double max);
	int set_x_final_ae(int constraint_index, const std::vector<double> x_index,
					   const std::vector<double> x_factor, double offset);
	int set_x_final_ae_offset(int constraint_index, double offset);
	int set_x_final_slack_ineq(int constraint_index, const std::vector<double> x_index,
							   const std::vector<double> x_factor, double min, double max);
	int set_x_final_slack_ineq_bound(int constraint_index, double min, double max);
	int set_delta_u_ae(unsigned int constraint_index, unsigned int control_index, double offset);
	int set_delta_u_ae_offset(unsigned int constraint_index, double offset);
	int set_delta_u_slack_ineq(unsigned int constraint_index, unsigned int u_index, double min, double max);
	int set_delta_u_slack_ineq_bound(unsigned int constraint_index, double min, double max);
	int set_delta_x_ae(unsigned int constraint_index, unsigned int x_index, double offset);
	int set_delta_x_ae_offset(unsigned int constraint_index, double offset);
	int set_delta_x_slack_ineq(unsigned int constraint_index, unsigned int x_index, double min, double max);
	int set_delta_x_slack_ineq_bound(unsigned int constraint_index, double min, double max);
	bool update_mpc(const Eigen::VectorXd &x0, const Eigen::VectorXd &u_fb,
					const Eigen::MatrixXd &x_ref, const Eigen::MatrixXd &u_ref, const Eigen::MatrixXd &distur_ref,
					const Eigen::MatrixXd &q, const Eigen::MatrixXd &r, const std::vector<double> &xu_ae_w = {},
					const std::vector<double> &xu_slack_ineq_w = {}, const std::vector<double> &x_final_ae_w = {},
					const std::vector<double> &x_final_slack_ineq_w = {}, const std::vector<double> &delta_u_ae_w = {},
					const std::vector<double> &delta_u_slack_ineq_w = {}, const std::vector<double> &delta_x_ae_w = {},
					const std::vector<double> &delta_x_slack_ineq_w = {});
	bool update_mpc(
		const Eigen::MatrixXd &q, const Eigen::MatrixXd &r,
		const Eigen::VectorXd &x0, const Eigen::VectorXd &u_fb, const Eigen::MatrixXd &x_ref,
		const Eigen::MatrixXd &distur_ref);

	bool update_mpc(
		const Eigen::MatrixXd &q, const Eigen::MatrixXd &r,
		const Eigen::VectorXd &x0, const Eigen::MatrixXd &x_ref,
		const Eigen::MatrixXd &distur_ref); // dis_sys_ptr  with disturbance
	// dis_sys_ptr  without disturbance
	bool update_mpc(
		const Eigen::MatrixXd &q, const Eigen::MatrixXd &r,
		const Eigen::VectorXd &x0, const Eigen::MatrixXd &x_ref);
	bool mpc_solve();
	//bool step(const Eigen::MatrixXd &x_ref, const Eigen::MatrixXd &x0,
	//const Eigen::MatrixXd &distur_ref = Eigen::MatrixXd::Zero(4, 0));

	Eigen::MatrixXd q();
	Eigen::MatrixXd r();
	std::vector<double> xu_ae_w();
	std::vector<double> xu_slack_ineq_w();
	std::vector<double> x_final_ae_w();
	std::vector<double> x_final_slack_ineq_w();
	std::vector<double> delta_u_ae_w();
	std::vector<double> delta_u_slack_ineq_w();
	std::vector<double> delta_x_ae_w();
	std::vector<double> delta_x_slack_ineq_w();

	bool set_x0(const Eigen::VectorXd &x0);
	bool set_u_minus1(const Eigen::VectorXd &u_minus1);
	bool set_distur_ref(const Eigen::MatrixXd &distur_ref);
	bool set_x_ref(const Eigen::MatrixXd &x_ref);
	bool set_u_ref(const Eigen::MatrixXd &u_ref);
	Eigen::VectorXd x0();
	Eigen::MatrixXd distur_ref();
	Eigen::MatrixXd x_ref();
	Eigen::MatrixXd u_ref();
	void clear_changed_flag();
	unsigned int state_dim();
	unsigned int control_dim();
	unsigned int distur_dim();
	unsigned int mpc_windows();
	std::shared_ptr<DiscreteLinearSys> dis_sys();
	Eigen::SparseMatrix<double> hessian();
	Eigen::VectorXd gradient();

	Eigen::VectorXd get_optimal_ctr();
	Eigen::SparseMatrix<double> constraint_matrix();
	Eigen::VectorXd lowerBound();
	Eigen::VectorXd upperBound();
	Eigen::VectorXd qp_solution();
	std::vector<double> get_state_vector(int index); //index from 0 to state_dim-1
	std::vector<double> get_control_vector(int index);
	Eigen::VectorXd get_state_one_shot(unsigned int time_index);
	Eigen::VectorXd get_contorl_one_shot(unsigned int time_index);
	Eigen::VectorXd optimal_ctr();

	//distur_ref cols eithr 1 or mpc_windows+1,1 means constant ref
	bool set_q(const Eigen::MatrixXd &q);
	bool set_r(const Eigen::MatrixXd &r);
	bool set_xu_ae_w(const std::vector<double> &xu_ae_w);
	bool set_xu_slack_ineq_w(const std::vector<double> &xu_slack_ineq_w);
	bool set_x_final_ae_w(const std::vector<double> &x_final_ae_w);
	bool set_x_final_slack_ineq_w(const std::vector<double> &x_final_slack_ineq_w);
	bool set_delta_u_ae_w(const std::vector<double> &delta_u_ae_w);
	bool set_delta_u_slack_ineq_w(const std::vector<double> &delta_u_slack_ineq_w);
	bool set_delta_x_ae_w(const std::vector<double> &delta_x_ae_w);
	bool set_delta_x_slack_ineq_w(const std::vector<double> &delta_x_slack_ineq_w);

protected:
	std::shared_ptr<DiscreteLinearSys> dis_sys_ptr_;
	unsigned int mpc_windows_;

	unsigned int max_xu_constraint_;
	unsigned int max_delta_x_;
	unsigned int max_delta_u_;
	unsigned int max_final_constraint_;

	unsigned int max_xu_ae_;
	unsigned int max_xu_slack_ineq_;
	unsigned int max_x_final_ae_;
	unsigned int max_x_final_slack_ineq_;
	unsigned int max_delta_u_ae_;
	unsigned int max_delta_u_slack_ineq_;
	unsigned int max_delta_x_ae_;
	unsigned int max_delta_x_slack_ineq_;

	std::vector<unsigned int> delta_u_u_index_;
	std::vector<double> delta_u_min_;
	std::vector<double> delta_u_max_;

	unsigned int state_dim_;
	unsigned int control_dim_;
	unsigned int distur_dim_;
	//input
	Eigen::VectorXd x0_;
	Eigen::VectorXd u_minus1_;
	Eigen::MatrixXd distur_ref_;
	bool is_distur_ref_changed_ = true;
	Eigen::MatrixXd x_ref_;
	bool is_x_ref_changed_ = true;
	Eigen::MatrixXd u_ref_;
	bool is_u_ref_changed_ = true;
	//weighting factor
	Eigen::MatrixXd q_;
	bool is_q_changed_ = true;
	Eigen::MatrixXd r_;
	bool is_r_changed_ = true;
	std::vector<double> xu_ae_w_;
	bool is_xu_ae_w_changed_ = true;
	std::vector<double> xu_slack_ineq_w_;
	bool is_xu_slack_ineq_w_changed_ = true;
	std::vector<double> x_final_ae_w_;
	bool is_x_final_ae_w_changed_ = true;
	std::vector<double> x_final_slack_ineq_w_;
	bool is_x_final_slack_ineq_w_changed_ = true;
	std::vector<double> delta_u_ae_w_;
	bool is_delta_u_ae_w_changed_ = true;
	std::vector<double> delta_u_slack_ineq_w_;
	bool is_delta_u_slack_ineq_w_changed_ = true;
	std::vector<double> delta_x_ae_w_;
	bool is_delta_x_ae_w_changed_ = true;
	std::vector<double> delta_x_slack_ineq_w_;
	bool is_delta_x_slack_ineq_w_changed_ = true;

	//anchor variable
	unsigned int u_1_row_start_index_;
	unsigned int xu_row_start_index_;
	unsigned int delta_x_row_start_index_;
	//unsigned int delta_u_1_row_start_index_;
	unsigned int delta_u_row_start_index_;
	unsigned int final_row_start_index_;
	unsigned int xu_ae_row_start_index_;
	unsigned int xu_slack_ineq_row_start_index_;
	unsigned int x_final_ae_row_start_index_;
	unsigned int x_final_slack_ineq_row_start_index_;
	unsigned int delta_u_ae_row_start_index_;
	unsigned int delta_u_slack_ineq_row_start_index_;
	unsigned int delta_x_ae_row_start_index_;
	unsigned int delta_x_slack_ineq_row_start_index_;

	unsigned int constriants_num_;

	unsigned int u_1_col_start_index_;
	unsigned int u_col_start_index_;
	unsigned int xu_ae_col_start_index_;
	unsigned int xu_slack_ineq_col_start_index_;
	unsigned int x_final_ae_col_start_index_;
	unsigned int x_final_slack_ineq_col_start_index_;
	unsigned int delta_u_ae_col_start_index_;
	unsigned int delta_u_slack_ineq_col_start_index_;
	unsigned int delta_x_ae_col_start_index_;
	unsigned int delta_x_slack_ineq_col_start_index_;

	unsigned int opti_var_num_;

	OsqpSolver::Solver solver_;
	Eigen::SparseMatrix<double> hessian_;
	Eigen::VectorXd gradient_;
	Eigen::SparseMatrix<double> constraint_matrix_;

	Eigen::VectorXd lowerBound_;
	Eigen::VectorXd upperBound_;
	Eigen::VectorXd qp_solution_;
	Eigen::VectorXd optimal_ctr_;
	bool is_solver_inited_ = false;

	void set_hessian(); //q r related
	//q x_ref related
	void set_gradient(); //x_ref cols eithr 1 or mpc_windows+1,1 means constant ref
	//sys distur_ref related
	void set_sf_constraint();
	bool init_osqp_solver();
};
} // namespace LinearMpc